package com.kqb.sgai.chatmemory;

import com.kqb.sgai.model.document.Msg;
import com.kqb.sgai.model.enums.MsgTypeEnum;
import com.kqb.sgai.service.MsgService;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.*;
import org.springframework.stereotype.Component;

import java.util.ArrayList;
import java.util.Date;
import java.util.List;

/**
 * @author 杉果树皮
 */
@Slf4j
@Component
public class MongodbChatMemory implements ChatMemory {

    @Resource
    private MsgService msgService;

    @Override
    public void add(String conversationId, List<Message> messages) {
        List<Msg> msgs = toMsg(messages, conversationId);
        msgService.saveMsg(msgs);
    }

    @Override
    public List<Message> get(String conversationId, int lastN) {
        List<Msg> msgs = msgService.findByConversationId(conversationId, lastN);
        return toMessage(msgs);
    }

    @Override
    public void clear(String conversationId) {
        msgService.deleteByConversationId(conversationId);
    }

    private List<Msg> toMsg(List<Message> messages, String conversationId) {
        if (messages == null) {
            return null;
        }
        ArrayList<Msg> msgs = new ArrayList<>();
        for (Message message : messages) {
            Msg msg = new Msg();
            msg.setConversationId(conversationId);
            msg.setText(message.getText());
            msg.setDate(new Date());
            msg.setType(MsgTypeEnum.getEnumByValue(message.getMessageType().getValue()));
            msgs.add(msg);
        }
        return msgs;
    }

    private List<Message> toMessage(List<Msg> msgs) {
        if (msgs == null) {
            return null;
        }
        ArrayList<Message> messages = new ArrayList<>();
        for (Msg msg : msgs) {
            switch (msg.getType()) {
                case USER -> messages.add(new UserMessage(msg.getText()));
                case ASSISTANT -> messages.add(new AssistantMessage(msg.getText()));
                case SYSTEM -> messages.add(new SystemMessage(msg.getText()));
                default -> log.error("不支持的消息类型");
            }
        }
        return messages;
    }
}
